# -*- coding: utf-8 -*-
import os

import numpy as np
import nibabel as nib

###Load one patient###
def returnPatients(path):
    if len(path) == 1:
        img = nib.load(path[0]+'/img.nii.gz')
        imgData = img.get_data()
        return imgData[np.newaxis,:]
    else:
        imgs = []
        for i in range(len(path)):
            img = nib.load(path[i]+'/img.nii.gz')
            imgs.append(img.get_data())
        return imgs


def preprocess(path):
    rawData = returnPatients(path)
    numOfPatiens = len(rawData)
    for i in range(numOfPatiens):
        train_x = rawData[i].copy()
        new_train_x = []
        for j in range(2, len(train_x) - 2):
            new_train_x.append(train_x[j-2:j+3])
    new_train_x = np.array(new_train_x)
    return new_train_x
